import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from mpl_toolkits.mplot3d import Axes3D  
import os
'''守恒格式测试程序。修改好了边界条件'''

'''functions'''
#sigmoid_lm1_r0,对x较小时为-1，较大时为0。用于构造平滑的阶跃函数。k越大越陡，c决定中心位置
def sigmoid_lm1_r0(x, c=0, k=5):  
    return -1 / (1 + np.exp(k * (x - c)))   



'''输入速度，计算涡度。采用中心差分方法。
函数本义是通用的，但此程序中由于没有更新u,v在中间区域的值，所以就在设初值时用了一次
注意，左侧指标是j，右侧指标是i
最上方：[-1, 1:-1]
'''
def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega


'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream(omega_int, dx, dy, tol=1e-5, psi_init=None):
    """
    Jacobi iteration
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    change = 1.
    iter = 0
    while change > tol:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)

        
        change = np.max(np.abs(psi_new - psi_init))  #change为迭代后矩阵元素变化最大值

        psi_init = psi_new
        iter += 1   
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    return psi_init, end_signal




#边界涡度处理，采用无滑移条件，顶盖速度恒为1，其余三壁速度恒为零.(太耍赖了，用psi和u,v)
def apply_bnd_omega_(psi, u, v, omega):
    #原来的左右壁面涡量条件
    omega[:, 0] = 2 / dx / dx * (psi[:, 0] - psi[:, 1]) - 2 / dx * v[:, 0]
    omega[:, -1] = 2 / dx / dx * (psi[:, -1] - psi[:, -2]) + 2 / dx * v[:, -1]

    # #左右边界周期条件
    # omega[:,0]=omega[:,-2]
    # omega[:,-1]=omega[:,1]

    #如果底部采用第二类边条件.这实际上是出口条件，不管了。
    #omega[0, 1:-1] =omega[1, 1:-1] 
    
    #如果底部仍用壁面边条件
    omega[0, 1:-1] = 2 / dy / dy * (psi[0, 1:-1] - psi[1, 1:-1]) + 2 / dy * u[0, 1:-1]
    
    #顶部采用的壁面涡量条件，这是顶部切向速度要求
    omega[-1, 1:-1] = 2 / dy / dy * (psi[-1, 1:-1] - psi[-2, 1:-1]) - 2 / dy * u[-1,1 :-1]

#根据流函数计算u,v
def psi_to_vel(psi,u,v,dx,dy):
    #内点
    u[1:-1,1:-1]=(psi[2:, 1:-1] - psi[:-2, 1:-1])/2/dy
    v[1:-1,1:-1]=-(psi[1:-1, 2:] - psi[1:-1, :-2])/2/dx
    #边界
    u[:,0]=0
    u[:,-1]=0
    u[-1, 1:-1] = u_up
    u[0, 1:-1] = 0

    v[:,0]=0
    v[:,-1]=0
    v[-1, 1:-1] = 0 
    v[0, 1:-1] = 0
    return u,v

def forward_omega_FTCS(psi, omega,u,v):
    #原来的非守恒形式：
    #uw_x = u[1:-1,1:-1] / 2 / dx  * (omega[1:-1, 2:] - omega[1:-1, :-2])
    #试着改成守恒形式吧！
    uw_x =  (u[1:-1, 2:]*omega[1:-1, 2:] -u[1:-1, :-2]*omega[1:-1, :-2])/ 2 / dx 
    vw_y =  (v[2:, 1:-1]*omega[2:, 1:-1] - v[:-2, 1:-1]*omega[:-2, 1:-1])/2 / dy
    diff_y = nu * (omega[2:, 1:-1] + omega[:-2, 1:-1] - 2 * omega[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (omega[1:-1, 2:] + omega[1:-1, :-2] - 2 * omega[1:-1, 1:-1]) / dx / dx
    return omega[1:-1, 1:-1] + dt * (-uw_x + -vw_y + diff_y + diff_x)

def diag_u(u,out_iter):
    
    subfolder='u_y_section' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    plt.plot(y,u[:,Nx//2])
    plt.xlabel('y')  
    plt.ylabel('u') 
    plt.title('Re={},t={}'.format(Re,t_diag))  
    # 定义文件名，包含idiag的值  
    filename_base = 'u_y_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def diag_uv_quiver(u,v,xx,yy,out_iter):
    
    subfolder='uv_quiver' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    # 使用quiver函数绘制矢量图  
    Q = plt.quiver(xx, yy, u, v, units='width', scale=10, color='b')  
    plt.xlabel('X')  
    plt.ylabel('Y')  
    plt.title('2D Vector Field , Re={},t={}'.format(Re,t_diag))  
    # 定义文件名，包含idiag的值  
    filename_base = 'uv_quiver_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

#设置参数
Nx = 64+2
Ny = 64+2
Lx = Ly = 1
u_up=1  #顶盖速度
x = np.linspace(0, Lx, Nx)  #linspace生成的点列是包含头和尾的。
dx = x[-1] - x[-2]  #网格长
y = np.linspace(0, Ly, Ny)
dy = y[-1] - y[-2]
nu = 0.001
dt = 0.001  #时间步长

ntime=4000  #时间步数
ndiag=50 #每隔多少时间输出一次
Re=1/nu
xx, yy = np.meshgrid(x, y)  
'''meshgrid生成的矩阵由Ny个长Nx的数列构成，矩阵形状为Ny*Nx。这在索引时造成的顺序是xx[j][i]，如果i是x向指标、j是y向指标的话。
这样，指标数是(i,j)的格点，索引为[j][i]，空间绝对位置(xx[j][i],yy [ j ][i])
'''

#初始条件设置
#设置初始流函数场
psi=np.zeros_like(yy)
# phi0_y=0.1*sigmoid_lm1_r0(y,c=Ly/2,k=10/Ly)
# for i in range(0,Nx):
#     for j in range(0,Ny):
#         psi[j][i]=phi0_y[j]

#设置初始速度场
u = np.zeros_like(xx)
v = np.zeros_like(yy)
#u,v=psi_to_vel(psi,u,v,dx,dy)   #用流函数设置
u[-1, 1:-1] = u_up    #最上方速度。这里设置速度时排除了首尾格点。numpy可以一个数组等于一个数，这样把这个数赋值给这个数组
# #初始均匀流
# for i in range(0,Nx):
#     for j in range(0,Ny):
#         u[j][i]=yy[j][i]*u_up

omega = np.zeros_like(xx)
omega[1:-1, 1:-1] = vel_to_omega(u, v, dx, dy)  #计算涡度场。依然排除首尾格点
psi, _ = compute_stream(omega[1:-1, 1:-1], dx, dy)  #omega只需要中间(Nx-2)*(Ny-2)小方块的信息
u,v=psi_to_vel(psi,u,v,dx,dy)
apply_bnd_omega_(psi, u, v, omega)  #设置边界涡量。

#进入主循环
end_signal = False  #终止判断条件。可以再加一个判断条件。
out_iter = 0    
#while end_signal==False and out_iter<1000:
while  out_iter<ntime:
    # if out_iter % ndiag==0:
    # #     diag_u(u,out_iter)    #截面速度分布
    #     diag_uv_quiver(u,v,xx,yy,out_iter)

    out_iter+=1
    omega[1:-1, 1:-1] = forward_omega_FTCS(psi, omega,u,v)  #时间向前步进。此时边界涡量会向中心扩散
    psi, end_signal = compute_stream(omega[1:-1, 1:-1], dx, dy, psi_init=psi, tol=1e-5)  #泊松方程计算流函数。算出的流函数将被用于下一次边界设置和时间步进。
    #需要新增代码：根据流函数场计算速度场。
    u,v=psi_to_vel(psi,u,v,dx,dy)
    apply_bnd_omega_(psi, u, v, omega)  #设置边界涡量。u,v仅用到边界值，其不在程序中改变。因此这里的u,v并不是全流场速度。

    
    
t_total=out_iter*dt #总时间

#绘制流函数二维颜色图。
flag_contourf=False
if flag_contourf:
    fig, ax = plt.subplots()
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(psi.min(), psi.max(), 50)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, psi, levels=levels, cmap=plt.get_cmap('Spectral'))  
    # 添加标题  
    ax.set_title('Re={},t={}'.format(Re,t_total), fontsize=16)  # 设置标题文本和字体大小  
    #添加colorbar
    cbar = fig.colorbar(cs)
    plt.show()

#绘制流函数三维图
    

flag_axe3D=False
if flag_axe3D:
    # 创建三维图形和坐标轴  
    fig = plt.figure()  
    ax = fig.add_subplot(111, projection='3d')  
    
    # 绘制三维曲面  
    ax.plot_surface(xx, yy, psi, cmap='viridis', linewidth=0)  # cmap参数用于设置颜色映射  
    
    # 添加标题和坐标轴标签  
    ax.set_title('3D Plot of psi')  
    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('Psi')  
    
    # 显示图形  
    plt.show()

#绘制截面速度图
flag_u_ysection=False
if flag_u_ysection:
    #增设画图部分。
    plt.plot(y,u[:,Nx//2])
    plt.title('Re={},t={}'.format(Re,t_total))  
    plt.show()

flag_quiver=True
if flag_quiver:
    t_diag=t_total
    plt.figure(figsize=(12, 8), dpi=100)  
    # 使用quiver函数绘制矢量图  
    Q = plt.quiver(xx, yy, u, v, units='width', scale=10, color='b')  
    plt.xlabel('X')  
    plt.ylabel('Y')  
    plt.title('2D Vector Field , Re={},t={}'.format(Re,t_diag))
    plt.show()